# -*- coding: utf-8 -*-

import ast

from ir.ir_ds import ITypeStruct, IValueStruct
from ir.ir_helper import Helper
from ir.ir_node import ITypeNode, IExpThis, IVarPort, IVarReg, IVarConst, IVarTemp, IVarParam, IVarNode
from ir.ir_statement import *

from .convert_exp import convert_exp, fetch_port_type
from .convert_type import convert_ast_to_type


def get_member_chain(ir_exp) -> list:
    # Get the parent member chain.
    if isinstance(ir_exp, IExpMember):
        return get_member_chain(ir_exp.parent) + [ir_exp.attr]
    elif isinstance(ir_exp, IExpThis):
        return ['this']
    elif isinstance(ir_exp, IVarReg):
        return ['this', 'reg', ir_exp.name]
    elif isinstance(ir_exp, IExpSubscript):
        # Append the current name to source chain.
        return get_member_chain(ir_exp.parent) + [ir_exp.index]
    raise Exception('Unknown parent {}'.format(type(ir_exp.parent)))


def parse_node_create(ir_call: IOpCall):
    # Check type.
    if not isinstance(ir_call.params[0], ITypeNode):
        raise Exception('Node initialize function should set Node type as the first parameter.')
    # Give back the Node type and parameters.
    return ir_call.params[0]


def update_this_table(ir_assign: IOpAssign, helper: Helper, this: ITypeNode):
    # Check for temp var assign.
    if isinstance(ir_assign.target, IVarTemp):
        # No need to update for temporary variable.
        return
    # Check whether the item is defined in node.
    if not isinstance(ir_assign.target, IExpMember):
        # If it is already a variable, just ignore it.
        if isinstance(ir_assign.target, IVar) or \
                isinstance(ir_assign.target, IExpSubscript):
            return
        raise Exception('{} is not a member type expression.'.format(ir_assign.target))
    # Get the chain of the item definition.
    target_chain = get_member_chain(ir_assign.target)
    # Basically, this chain should be 3 items.
    if len(target_chain) > 3:
        # For more than 3 items, it must have some further extraction, ignore.
        return
    if len(target_chain) != 3:
        raise Exception('Unexpected target chain detected: {}'.format(target_chain))
    # The second string defines the target dictionary in this.
    dict_name = target_chain[1]
    this_dict = getattr(this, dict_name + 's')
    # The third string defines the variable name in this.
    var_name = target_chain[2]
    # If variable name already exist in dictionary.
    if var_name in this_dict:
        return
    # Now the variable is ensured not in dictionary, based on the value, to
    # construct the data. Check for the types of value by checking the dict
    # name.
    if dict_name == 'port':
        # So now checking the assignment value type.
        ir_exp = ir_assign.exp
        if isinstance(ir_exp, IOpCall):
            # Fetch the port type.
            ir_port_type = fetch_port_type(ir_exp, helper)
            # Now construct the data.
            ir_var = IVarPort(var_name, ir_port_type, ir_assign.exp)
            # Check the generic assignment.
            if hasattr(ir_port_type, 'generic'):
                # Copy the generic types.
                ir_var.generic = ir_port_type.generic
        elif isinstance(ir_exp, IValueListComp):
            # Convert the port type from the function call.
            ir_port_type = fetch_port_type(ir_exp.exp, helper)
            # Update the list definition.
            ir_exp.type.managed_type = ir_port_type
            # So a port is assign as a list, that's simple to define the port.
            ir_var = IVarPort(var_name, ir_exp.type, IExpValue(ir_exp))
        elif isinstance(ir_exp, IExpMember):
            # Extract the port type first.
            ir_port_type = fetch_port_type(ir_exp, helper)
            # Now construct the port.
            ir_var = IVarPort(var_name, ir_port_type, ir_exp)
        else:
            raise Exception('Unimplemented port assignment exp type {}.'.format(type(ir_exp)))
    elif dict_name == 'reg':
        # Construct the reg variable.
        ir_var = IVarReg(var_name, convert_ast_to_type(ir_assign.exp, helper))
    elif dict_name == 'const':
        # Construct the const variable.
        ir_var = IVarConst(var_name, convert_ast_to_type(ir_assign.exp, helper), ir_assign.exp)
    elif dict_name == 'node':
        ir_exp = ir_assign.exp
        if isinstance(ir_exp, IOpCall):
            # Directly construct the node type and parameter by parsing the call.
            node_type = parse_node_create(ir_exp)
        elif isinstance(ir_exp, IValueListComp):
            # Parse the managed node type.
            item_node_type = parse_node_create(ir_exp.exp)
            # Update the managed item type.
            ir_exp.type.managed_type = item_node_type
            # Now the node type is the expression type.
            node_type = ir_exp.type
        else:
            raise Exception('Unknown node construction type: {}.'.format(type(ir_exp)))
        # Construct the node variable.
        ir_var = IVarNode(var_name, node_type)
    else:
        raise Exception('Unimplemented assignment checking type {}.'.format(dict_name))
    # Add the name to dictionary.
    this_dict[var_name] = ir_var
    # Update the assign target.
    ir_assign.target = ir_var


def update_this_table_from_list(assign_result: list, helper: Helper, this: ITypeNode):
    # Loop for all the assignment ir in the result
    for ir_assign in assign_result:
        update_this_table(ir_assign, helper, this)


def convert_assign(ast_assign: ast.Assign, helper: Helper, this: ITypeNode) -> list:
    # Get the target of the value.
    assert len(ast_assign.targets) == 1
    # Get the target item.
    ass_target = convert_exp(ast_assign.targets[0], helper, this)
    # Convert the value part.
    ass_value = convert_exp(ast_assign.value, helper, this)
    # It might contain multiple assignment, convert those assignment into
    # multiple assignment.
    if isinstance(ass_value, IOpCall):
        # The return value of the function call should be able to extracted.
        func_def = ass_value.funcdef
        ret_type = func_def.ret_type
        # Check the return type.
        if isinstance(ret_type, ITypeStruct) and ret_type.extract:
            # Okay, we could extract value from the return type.
            # First, we need a temp var to hold it.
            ir_temp = IVarTemp('ret_{}'.format(id(ast_assign)), ret_type)
            assign_result = [IOpAssign(ir_temp, ass_value)]
            # Now we need to extract the result.
            if isinstance(ass_target, IExpSet):
                # Construct the variable name and type list.
                var_nts = ass_target.data
            else:
                # Treat the target as a single item.
                var_nts = [ass_target]
            # Loop for all the type data.
            for item, item_name, item_type in zip(var_nts, ret_type.member, ret_type.member_type):
                # Ignore the void assignment
                if item == IExpVoid:
                    continue
                # Construct the member exp.
                assign_result.append(IOpAssign(item, IExpMember(ir_temp, item_name, item_type)))
        else:
            # Just treat as single assignment.
            assign_result = [IOpAssign(ass_target, ass_value)]
    else:
        assign_result = [IOpAssign(ass_target, ass_value)]
    # Base on the assign variable, check the definition table.
    update_this_table_from_list(assign_result, helper, this)
    # Give back the assignment result.
    return assign_result


def convert_annassign(ast_annassign: ast.AnnAssign, helper: Helper, this: ITypeNode):
    # Basically the annotation assign is a class definition for a type of variable.
    ass_target = convert_exp(ast_annassign.target, helper, this)
    # Convert the annotation.
    ass_type = convert_ast_to_type(ast_annassign.annotation, helper)
    # Check out the target.
    if isinstance(ass_target, IVarParam):
        # A parameter type define is shown.
        ass_target.type = ass_type
        # Add the defnition to this.
        this.params[ass_target.name] = ass_target
        return []
    raise Exception('Unknown target ir {} is detected.'.format(ass_target))


def convert_augassign(ast_augassign: ast.AugAssign, helper: Helper, this: ITypeNode):
    # Get the target item.
    ass_target = convert_exp(ast_augassign.target, helper, this)
    # Convert the value part.
    ass_value = convert_exp(ast_augassign.value, helper, this)
    # Based on the operation, define the op.
    if isinstance(ast_augassign.op, ast.Add):
        op = IOpAugAssign.ADD
    elif isinstance(ast_augassign.op, ast.Sub):
        op = IOpAugAssign.SUB
    elif isinstance(ast_augassign.op, ast.Mult):
        op = IOpAugAssign.MULT
    elif isinstance(ast_augassign.op, ast.Div):
        op = IOpAugAssign.DIV
    elif isinstance(ast_augassign.op, ast.Mod):
        op = IOpAugAssign.MOD
    elif isinstance(ast_augassign.op, ast.BitOr):
        op = IOpAugAssign.BOR
    elif isinstance(ast_augassign.op, ast.BitAnd):
        op = IOpAugAssign.BAND
    elif isinstance(ast_augassign.op, ast.BitXor):
        op = IOpAugAssign.BXOR
    elif isinstance(ast_augassign.op, ast.LShift):
        op = IOpAugAssign.BLSHIFT
    elif isinstance(ast_augassign.op, ast.RShift):
        op = IOpAugAssign.BLSHIFT
    else:
        raise Exception('Unknown AugAssign operation: {}.'.format(ast_augassign.op))
    # If you could do AugAssign, which means the target must exist before.
    return [IOpAugAssign(ass_target, ass_value, op)]